package com.maomao.es.vector;

import org.elasticsearch.core.Map;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.script.Script;
import org.elasticsearch.script.ScriptType;
import org.springframework.data.domain.Page;
import org.springframework.data.domain.Pageable;
import org.springframework.data.elasticsearch.annotations.Query;
import org.springframework.data.elasticsearch.core.query.NativeSearchQuery;
import org.springframework.data.elasticsearch.core.query.NativeSearchQueryBuilder;
import org.springframework.data.elasticsearch.repository.ElasticsearchRepository;

/**
 * @author mao
 * @date 2025-05-30- 14:53
 * @descript
 */
public interface VectorDocumentRepository extends ElasticsearchRepository<VectorDocument, String> {

    // 使用script_score进行向量相似度搜索
    // @Query("{\"script_score\": {\"query\": {\"match_all\": {}}, \"script\": {\"source\": \"cosineSimilarity(params.query_vector, 'contentVector') + 1.0\", \"params\": {\"query_vector\": ?0}}}}")
    @Query("{\n" +
            "    \"script_score\": {\n" +
            "      \"query\": {\"match_all\": {}}, \n" +
            "      \"script\": {\n" +
            "        \"source\": \"cosineSimilarity(params.query_vector, 'titleVector') + 1.0\", \n" +
            "        \"params\": {\n" +
            "          \"query_vector\": ?0\n" +
            "        }\n" +
            "      }\n" +
            "    }\n" +
            "  }")
    Page<VectorDocument> findSimilarByContentVector(float[] queryVector, Pageable pageable);

    // 混合搜索：结合关键词和向量
    @Query("{\"query\": {\"bool\": {\"should\": [{\"match\": {\"content\": \"?0\"}}, {\"script_score\": {\"query\": {\"match_all\": {}}, \"script\": {\"source\": \"cosineSimilarity(params.query_vector, 'contentVector') + 1.0\", \"params\": {\"query_vector\": ?1}}}}]}}}")
    Page<VectorDocument> hybridSearch(String keyword, float[] queryVector, Pageable pageable);

    // 带权重的混合搜索
//    default Page<VectorDocument> weightedHybridSearch(
//            String keyword,
//            float[] queryVector,
//            float keywordWeight,
//            float vectorWeight,
//            Pageable pageable) {
//
//        String script = String.format(
//                "(_score * %f) + (cosineSimilarity(params.query_vector, 'contentVector') * %f",
//                keywordWeight,
//                vectorWeight
//        );
//
//        NativeSearchQuery query = new NativeSearchQueryBuilder()
//                .withQuery(QueryBuilders.boolQuery()
//                        .should(QueryBuilders.matchQuery("content", keyword).boost(keywordWeight))
//                        .should(QueryBuilders.scriptScoreQuery(
//                                QueryBuilders.matchAllQuery(),
//                                new Script(ScriptType.INLINE, "painless", script,
//                                        Map.of("query_vector", queryVector))
//                        ).boost(vectorWeight))
//                )
//                .withPageable(pageable)
//                .build();
//
//        return repository.search(query);
//    }
}
